# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import concurrent.futures
import datetime
import logging
import multiprocessing
import os
import matplotlib
matplotlib.use("Agg")

from absl import app, flags

from coin.base.datetime_util import to_datetime
from coin.strategy.tool.monitor.log_parser_runner import check_utc, run_config_v1
from coin.strategy.tool.monitor import config_v1_1, config_v1_2, config_v2
from coin.strategy.tool.monitor.plot import (setup_plt,
                                             get_plotter,
                                             refine_balance,
                                             refine_timeseries)


def plot_balance(results, begin_time, end_time, mark_time, filename, axis_interval_hour=None):
  assert begin_time < end_time
  axis_interval_hour = max(int((end_time - begin_time).total_seconds() / 3600 / 10), 1)
  with get_plotter(filename,
                   begin_time,
                   end_time,
                   plot_secondary=True,
                   axis_interval_hour=axis_interval_hour) as plots:
    for result in results:
      try:
        ret_bps = "%.2f" % ((result['pnl_mark'] / result['turnover_mark']) * 1e4)
        pnl_mark = "%.2f" % result['pnl_mark']
        turnover_mark = "%.2f" % result['turnover_mark']
        title = f'pnl {pnl_mark} / tamt {turnover_mark} = retbps {ret_bps}'
        result['name'] = '%s %s' % (result['name'], title)
      except Exception:
        pass

      plots[0](*refine_balance(result, begin_time, end_time, mark_time))
      plots[1](*refine_timeseries(result, begin_time, end_time, 'position_dump'))


def query_balance_from_log(from_dt, to_dt):
  configs = config_v1_1.STRATEGIES + config_v1_2.STRATEGIES + config_v2.LOG_FILES
  futures = []
  max_workers = multiprocessing.cpu_count()
  with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) \
       as executor:
    for config in configs:
      if not isinstance(config, dict):
        continue
      fut = executor.submit(run_config_v1, config, to_dt, to_dt - from_dt)
      futures.append(fut)
  return [f.result() for f in futures]


def _aggregate_by_key(iterable, key):
  res = {}
  for elem in iterable:
    k = elem[key]
    if k not in res:
      res[k] = []
    res[k].append(elem)
  return res


def main(argv):
  FLAGS = flags.FLAGS

  assert check_utc(), 'Invalid timezone'

  from_dt = datetime.datetime.strptime(FLAGS.from_date, '%Y%m%d')
  if FLAGS.to_date is None:
    to_dt = to_datetime(datetime.date.today())
  else:
    to_dt = datetime.datetime.strptime(FLAGS.to_date, '%Y%m%d')
  filename_suffix = '%s-%s' % (from_dt.strftime('%Y%m%d'), to_dt.strftime('%Y%m%d'))
  print('Dump period: %s - %s' % (from_dt.strftime('%Y%m%d'), to_dt.strftime('%Y%m%d')))

  to_dt += datetime.timedelta(days=1)
  results = query_balance_from_log(from_dt, to_dt)
  results_by_group = _aggregate_by_key(results, 'group')

  plot_dir = os.path.normpath(os.path.expanduser(FLAGS.plot_dir))
  os.makedirs(plot_dir, exist_ok=True)

  for group, res_list in results_by_group.items():
    filename = os.path.join(plot_dir, '%s.%s.png' % (group, filename_suffix))
    plot_balance(res_list, from_dt, to_dt, from_dt, filename)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  flags.DEFINE_string('plot_dir', 'plot_out', 'Plot directory')
  flags.DEFINE_string('from_date', None, 'From date')
  flags.DEFINE_string('to_date', None, 'To date. If not specified, today is used')
  app.run(main)
